# Code for radial basis functions.

import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk

from typing import Optional, Union, Tuple, List
from jax import nn as nn

_RBF_BY_KEY = dict()

def _rbf_register(*aliases):
    """Return the alias register"""
    def alias_reg(cls):
        name = cls.__name__.lower()
        if name not in _RBF_BY_KEY:
            _RBF_BY_KEY[name] = cls
        for alias in aliases:
            if alias not in _RBF_BY_KEY:
                _RBF_BY_KEY[alias] = cls
        return cls
    return alias_reg

class RadialBasisFunctions(hk.Module):
    r"""Network for radial basis functions.

    Args:
        r_max (float):          Maximum distance. Default: 1.0 nm.
        r_min (float):          Minimum distance. Default: 0.0 nm.
        sigma (float):          Sigma. Default: 0.1 nm.
        num_basis (int):        Number of basis functions. Default: None.
        rescale (bool):         Whether to rescale the output of RBF from -1 to 1. Default: False
        clip_distance (bool):   Whether to clip the value of distance. Default: False
        hyper_param (dict):     Hyperparameter. Default: None
    
    """

    def __init__(self,
                 fp_type = jnp.float32,
                 r_max: float = 1.0,
                 r_min: float = 0.0,
                 num_basis: Optional[int] = None,
                 clip_distance: bool = False,
                 name: str = "radial_basis_functions",
                 ):
        super().__init__(name=name)

        self.fp_type = fp_type
        self.num_basis = num_basis
        self.r_max = r_max
        self.r_min = r_min
        self.clip_distance = clip_distance
    
    def __call__(self, distance: jnp.ndarray) -> jnp.ndarray:
        r"""Compute the radial basis functions.

        Args: 
            distance (Array):                 Distance matrix. Shape: (A, A).
        Returns: 
            radial basis embedding (Array):   Embedding of distance matrix. Shape: (A, A, num_basis).

        """
        raise NotImplementedError

    def __str__(self):
        return 'RadicalBasisFunctions<>'  
